-- re2hs $INPUT -o $OUTPUT
{-# OPTIONS_GHC -Wno-unused-record-wildcards #-}
{-# LANGUAGE OverloadedStrings #-}

import Control.Monad
import qualified Data.ByteString as BS
import GHC.IO.Handle
import System.Directory
import System.IO

chunk_size :: Int
chunk_size = 4096

data State = State {
    _file :: !Handle,
    _yyinput :: !BS.ByteString,
    _yycursor :: !Int,
    _yylimit :: !Int,
    _token :: !Int,
    _eof :: !Bool,
    _count :: !Int
}

%{
    re2c:YYFN = ["lexer;IO Int", "State{..};State;!State{..}"];
    re2c:YYPEEK = "BS.index";

    // We have to turn off autogenerated YFILL check and write it manually as part of YYFILL
    // implementation, so that we can propagate the updated state out of it.
    re2c:yyfill:check = 0;
    re2c:YYFILL = "State{..} <- fill State{..} @@";
    re2c:monadic = 1;

    str = ['] ([^'\\] | [\\][^])* ['];

    *      { return (-1) }
    [\x00] { return $ if _yycursor == _yylimit - yymaxfill + 1 then _count else (-1) }
    str    { lexer State{_token = _yycursor, _count = _count + 1, ..} }
    [ ]+   { lexer State{_token = _yycursor, ..} }
%}

yymaxfill :: Int
%{max %}

fill :: State -> Int -> IO State
fill !st@State{..} !need =
    if _yylimit - _yycursor >= need then
        return st
    else case _eof of
        True -> error "fill failed"
        False -> do
            -- Discard everything up to the current token, cut off terminating null,
            -- read new chunk from file and reappend terminating null at the end.
            chunk <- BS.hGet _file chunk_size
            let !eof = BS.length chunk < need -- end of file ?
            let !buf = BS.concat [
                    BS.drop _token _yyinput,
                    chunk,
                    if eof then (BS.replicate yymaxfill 0) else BS.empty]
            return State {
                _yyinput = buf,
                _yycursor = _yycursor - _token,
                _yylimit = BS.length buf,
                _token = 0,
                _eof = eof,
                ..}

main :: IO ()
main = do
    let fname = "input"

    -- Prepare input file.
    BS.writeFile fname $ BS.concat ["'qu\0tes' 'are' 'fine: \\'' " | _ <- [1..chunk_size]]
    let expect = 3 * chunk_size -- the total number of strings in file

    -- Run lexer on the prepared file.
    fh <- openFile fname ReadMode
    let st = State {
        _file = fh,
        _yyinput = BS.empty,
        _yycursor = 0,
        _token = 0,
        _yylimit = 0,
        _eof = False,
        _count = 0
    }
    result <- lexer st
    hClose fh

    -- Cleanup.
    removeFile fname

    -- Check result.
    when (result /= expect) $ error $ "expected " ++ show expect ++ ", got " ++ show result
    return ()
